# coding=utf-8
from pyspark.sql import SparkSession
from pyspark.sql.context import SQLContext

from tiancheng.base.base_helper import *
from pyspark.mllib.tree import GradientBoostedTrees, GradientBoostedTreesModel
from pyspark.mllib.util import MLUtils

# conf = SparkConf().setAppName('Gradient Boosted Tree Classification').setMaster('local[4]')
spark = SparkSession.builder.master('local[4]') \
        .appName("local[4]").getOrCreate()
sc = spark.sparkContext
# sqlContext = SQLContext(sparkContext=sc)
op_df = spark.read.csv(operation_train_new, header=True)
op_df.show()
# load and parse data file
data = MLUtils.loadLibSVMFile(sc, operation_train_new)

# split the data into training and test
trainingData, test = data.randomSplit([0.7, 0.3])

# train a gradient boost tree model
model = GradientBoostedTrees.trainClassifier(trainingData, categoricalFeaturesInfo={},
                                             numIterations=3)

# evaluate model on test instances and compute test error
predictions = model.predict(test.map(lambda x : x.features))
labelAndPredictions = test.map(lambda lp : lp.label).zip(predictions)
testErr = labelAndPredictions.filter(lambda v : v[0] != v[1]).count()/float(test.count())
print('test error :' + str(testErr))
print('learned classification GBT model :')
print(model.toDebugString)

# save and load
model.save(sc, './model/myGradientBoostingClassificationModel')
sameModel = GradientBoostedTreesModel.load(sc, './model/myGradientBoostingClassificationModel')

sc.stop()